import numpy as np
from typing import Callable, Dict, List, Tuple

def generate_nguyen_data(func: Callable, n_samples: int = 1000, noise: float = 0.1) -> Tuple[np.ndarray, np.ndarray]:
    """生成Nguyen方程的数据"""
    # 生成随机输入数据
    X = np.random.uniform(-1, 1, n_samples)
    # 计算目标值
    y = func(X)
    # 添加噪声
    y += np.random.normal(0, noise)
    y_ = np.array([])
    X_ = np.array([])
    for i, yitem in enumerate(y):
        if not np.isnan(yitem):
            y_ = np.append(y_, yitem)
            X_ = np.append(X_, X[i])
    return X_, y_

# Nguyen方程集合
NGUYEN_EQUATIONS: Dict[str, Callable] = {
    # 基础多项式方程
    "Nguyen-1": lambda x: x**3 + x**2 + x,  # x^3 + x^2 + x
    "Nguyen-2": lambda x: x**4 + x**3 + x**2 + x,  # x^4 + x^3 + x^2 + x
    "Nguyen-3": lambda x: x**5 + x**4 + x**3 + x**2 + x,  # x^5 + x^4 + x^3 + x^2 + x
    "Nguyen-4": lambda x: x**6 + x**5 + x**4 + x**3 + x**2 + x,  # x^6 + x^5 + x^4 + x^3 + x^2 + x
    "Nguyen-5": lambda x: np.sin(x**2) * np.cos(x) - 1,  # sin(x^2)cos(x) - 1
    
    # 三角函数方程
    "Nguyen-6": lambda x: np.sin(x) + np.sin(x + x**2),  # sin(x) + sin(x + x^2)
    "Nguyen-7": lambda x: np.log(x + 1) + np.log(x**2 + 1),  # log(x + 1) + log(x^2 + 1)
    "Nguyen-8": lambda x: np.sqrt(x),  # sqrt(x)
    "Nguyen-9": lambda x: np.sin(x) + np.sin(x**2),  # sin(x) + sin(x^2)
    "Nguyen-10": lambda x: 2 * np.sin(x) * np.cos(x),  # 2sin(x)cos(x)
    
    # 指数和对数方程
    "Nguyen-11": lambda x: x**x,  # x^x
    "Nguyen-12": lambda x: x**4 - x**3 + x**2/2 - x,  # x^4 - x^3 + x^2/2 - x
    "Nguyen-13": lambda x: 6 * np.sin(x) * np.cos(x),  # 6sin(x)cos(x)
    "Nguyen-14": lambda x: 8 / (2 + x**2),  # 8/(2 + x^2)
    "Nguyen-15": lambda x: np.sqrt(x**2 + 1),  # sqrt(x^2 + 1)
    
    # 复杂组合方程
    "Nguyen-16": lambda x: np.sin(x) + np.sin(x**2 + x),  # sin(x) + sin(x^2 + x)
    "Nguyen-17": lambda x: np.sin(x) + np.sin(x**2) + np.sin(x**3),  # sin(x) + sin(x^2) + sin(x^3)
}

def get_nguyen_equation(name: str) -> Tuple[Callable, str]:
    """获取指定的Nguyen方程及其描述"""
    if name not in NGUYEN_EQUATIONS:
        raise ValueError(f"未知的Nguyen方程: {name}")
    return NGUYEN_EQUATIONS[name], name

def get_all_nguyen_equations() -> List[Tuple[Callable, str]]:
    """获取所有Nguyen方程"""
    return [(func, name) for name, func in NGUYEN_EQUATIONS.items()]

def main():
    """测试Nguyen方程"""
    # 测试所有方程
    for name, func in NGUYEN_EQUATIONS.items():
        X, y = generate_nguyen_data(func, n_samples=100)
        print(f"\n测试方程 {name}:")
        print(f"输入形状: {X.shape}")
        print(f"输出形状: {y.shape}")
        # print(f"示例输入: {X}")
        # print(f"示例输出: {y}")

if __name__ == "__main__":
    main() 